#include "dag.h"
#include "dag_loop_analyzer.h"
#include "dag_resource_analyzer.h"

#include <list>
#include <set>
#include <vector>
#include <map>
//#include <iostream>

using namespace dag;
using namespace std;

//void print_set(const set<size_t>& s) {
//	set<size_t>::const_iterator siter;
//	for (siter = s.begin(); siter != s.end(); siter++) {
//		cout << *siter << "\t";
//	}
//	cout << endl;
//}

static void inputEdgesToSpec(const vector<Edge*>& edgeVec,
							 FrameSpec& spec,
							 list<GraphError>& errors);

static void outputEdgesToSpec(const vector<Edge*>& edgeVec,
							  FrameSpec& spec,
							  list<GraphError>& errors);

template<bool _treat_as_input_edges>
static void createSpecFromEdges(const vector<Edge*>& edgeVec,
								FrameSpec& spec,
								list<GraphError>& errors);

static bool startsAtZeroAndHasContinuity(const set<size_t>& sset);
static bool setOverlap(const set<size_t>& set1, const set<size_t>& set2);

template<typename _value, typename _key>
static const vector<_value> getValueVector(const map<_key, _value>& kvMap);
template<typename _value, typename _key>
static const set<_key> getKeySet(const map<_key, _value>& kvMap);
template<typename _value, typename _key>
static const _value getValue(const pair<_key, _value>& kvPair);
template<typename _value, typename _key>
static const _key getKey(const pair<_key, _value>& kvPair);



void dag::validate(const Graph& graph, 
					const FrameSpec& expectedInputSpec, 
					FrameSpec& outputSpec,
					list<GraphError>& errors) {
	list<set<const Node*> > loops;

	/* validate that no loops exist */
	find_loops(graph, loops);
	list<set<const Node*> >::const_iterator loopIter;
	for (loopIter = loops.begin(); loopIter != loops.end(); loopIter++) {
		errors.push_back(LoopError(*loopIter));
	}
	
	/* validate that all nodes are connected to something */
	const set<Node*>& nodes = graph.getNodes();
	set<Node*>::const_iterator nodeIter;
	for (nodeIter = nodes.begin(); nodeIter != nodes.end(); nodeIter++) {
		if (*nodeIter != graph.getHeadNode() || 
			*nodeIter != graph.getTailNode()) {
			if (graph.getOutputEdges(*nodeIter).size() == 0) {
				if (graph.getInputEdges(*nodeIter).size() == 0) {
					errors.push_back(UnconnectedNode(*nodeIter));
				}
			}
		}
	}
	
	/* validate edges and paths */
	validate_links(graph,
				   expectedInputSpec,
				   outputSpec,
				   errors);
	
	/* validate resources for learning. */
	validate_resources(graph, errors);
}

void dag::validate_resources(const Graph& graph, list<GraphError>& errors) {
	ResourceAnalyzer ra;
	ra.checkResources(graph, errors);
}

void dag::validate_links(const Graph& graph,
								const FrameSpec& expectedInputSpec,
								FrameSpec& outputSpec,
								list<GraphError>& errors) {
	
	if (!initialize(graph, expectedInputSpec, outputSpec, errors)) {
		return;
	}
	
	/* check that for each node, the entire node's output frame is mapped to 
	 * exactly one other node input */
	set<Node*>::const_iterator nodeIter;
	const set<Node*>& nodes = graph.getNodes();
	for (nodeIter = nodes.begin(); nodeIter != nodes.end(); nodeIter++) {
		if (*nodeIter != graph.getTailNode()) {
			FrameSpec spec;
			outputEdgesToSpec(graph.getOutputEdges(*nodeIter), spec, errors);
			if (spec != (*nodeIter)->outputSpec()) {
				errors.push_back(HangingConnection(*nodeIter,
												   graph.getOutputEdges(*nodeIter)));
			}
		}
	}	
}

bool dag::initialize(const Graph& graph,
					  const FrameSpec& expectedInputSpec,
					  FrameSpec& outputSpec,
					  list<GraphError>& errors) {
	map<const Node*, FrameSpec> nodeInputSpecs;
	list<GraphError> initErrors;
	nodeInputSpecs.clear();
	outputSpec = FrameSpec();
	
	/* go through each node in the graph and determine what the input spec would
	 * be based off of the edges leading into the node */
	const set<Node*>& nodes = graph.getNodes();
	set<Node*>::const_iterator nodeIter;
	for (nodeIter = nodes.begin(); nodeIter != nodes.end(); nodeIter++) {
		if (*nodeIter == graph.getHeadNode()) {
			/* don't process head node here since it doesn't have any input
			 * edges. */
			continue;
		}
		
		FrameSpec spec;
		inputEdgesToSpec(graph.getInputEdges(*nodeIter), 
						 spec,
						 initErrors);
		nodeInputSpecs[*nodeIter] = spec;
	}
	
	nodeInputSpecs[graph.getHeadNode()] = expectedInputSpec;
	
	/* call init on all nodes using their expected input.  Don't perform this
	 * if there were previous errors as a node could throw an error when
	 * initializing it with a faulty frame spec */
	if (initErrors.size() == 0) {
		for (nodeIter = nodes.begin(); nodeIter != nodes.end(); nodeIter++) {
			(*nodeIter)->init(nodeInputSpecs[*nodeIter]);
		}
		outputSpec = graph.getTailNode()->outputSpec();
		/* TODO: test for hanging connection here */
	}
	
	/* append errors from this module to the error list passed in */
	errors.insert(errors.end(), initErrors.begin(), initErrors.end());
	
	return (initErrors.size() == 0);
	
}

void dag::find_loops(const Graph& graph, 
					  list<set<const Node*> >& loops) {
	LoopAnalyzer loopAnalyzer;
	loopAnalyzer.findLoops(graph, loops);
}

NodeTraits dag::get_node_set_traits(const set<Node*>& nodes) {
	/* traits are determined from head node, tail node and any connected nodes */
	NodeTraits nodeTraits;
	nodeTraits.canReverse = true;
	nodeTraits.resources = LearnResource_None;
	
	/* check all nodes */
	set<Node*>::const_iterator nodeIter;
	for (nodeIter = nodes.begin();
		 nodeIter != nodes.end();
		 nodeIter++) {
		if (*nodeIter != NULL) {
			const NodeTraits& subNodeTraits = (*nodeIter)->traits();
			
			nodeTraits.canReverse &= subNodeTraits.canReverse;
			
			LearnResource& set_lr = nodeTraits.resources;
			const LearnResource& node_lr = subNodeTraits.resources;
			switch (set_lr) {
				case LearnResource_None:
					if (node_lr != LearnResource_None) {
						set_lr = node_lr;
					}
					break;
				case LearnResource_Inputs:
					if (node_lr == LearnResource_Outputs ||
						node_lr == LearnResource_InputsOutputs) {
						set_lr = LearnResource_InputsOutputs;
					}
					break;
				case LearnResource_Outputs:
					if (node_lr == LearnResource_Inputs ||
						node_lr == LearnResource_InputsOutputs) {
						set_lr = LearnResource_InputsOutputs;
					}
					break;
				case LearnResource_InputsOutputs:
					break;
			}
		}
	}
	return nodeTraits;
}

void inputEdgesToSpec(const vector<Edge*>& edgeVec,
								   FrameSpec& spec,
								   list<GraphError>& errors) {
	createSpecFromEdges<true>(edgeVec, spec, errors);
}

void outputEdgesToSpec(const vector<Edge*>& edgeVec,
									FrameSpec& spec,
									list<GraphError>& errors) {
	createSpecFromEdges<false>(edgeVec, spec, errors);
}

template<bool _treat_as_input_edges>
void createSpecFromEdges(const vector<Edge*>& edgeVec,
						 FrameSpec& spec,
						 list<GraphError>& errors) {
	spec = FrameSpec();
	
	/* go through all edges and determine the set of things connected.  During
	 * the process we can validate whether links overlap and return an
	 * error. */
	set<size_t> bConn, iConn, fConn, dConn;
	for (size_t i = 0; i < edgeVec.size(); i++) {
		/* these links pertain to this particular edge */
		set<size_t> bEConn, iEConn, fEConn, dEConn;
		const FrameLink& link = edgeVec[i]->link;
		
		if (_treat_as_input_edges) {
			bEConn = link.bLink.reverseSet();
			iEConn = link.iLink.reverseSet();
			fEConn = link.fLink.reverseSet();
			dEConn = link.dLink.reverseSet();
		}
		else {
			bEConn = link.bLink.forwardSet();
			iEConn = link.iLink.forwardSet();
			fEConn = link.fLink.forwardSet();
			dEConn = link.dLink.forwardSet();			
		}
		
		/* ensure that there is no overlap between these links and 
		 * links from other edges */
		if (setOverlap(bConn, bEConn) ||
			setOverlap(iConn, iEConn) ||
			setOverlap(fConn, fEConn) ||
			setOverlap(dConn, dEConn)) {
			
			/* Can't create spec.  Add error and exit */
			errors.push_back(AmbiguousConnection(edgeVec));
			return;
		}
		
		/* add to links */
		bConn.insert(bEConn.begin(), bEConn.end());
		iConn.insert(iEConn.begin(), iEConn.end());
		fConn.insert(fEConn.begin(), fEConn.end());
		dConn.insert(dEConn.begin(), dEConn.end());
	}
	
	/* validate the links go from 0 to N without missing any values */
	if (!(startsAtZeroAndHasContinuity(bConn) &&
		  startsAtZeroAndHasContinuity(iConn) &&
		  startsAtZeroAndHasContinuity(fConn) &&
		  startsAtZeroAndHasContinuity(dConn))) {
		
		const Node* node = NULL;
		if (edgeVec.size() > 0) {
			if (_treat_as_input_edges) {
				node = edgeVec[0]->to;
			}
			else {
				node = edgeVec[0]->from;
			}
		}		
		/* Can't create spec.  Add error and exit */
		errors.push_back(HangingConnection(node, edgeVec));
	}
	
	spec.bSize = bConn.size();
	spec.iSize = iConn.size();
	spec.fSize = fConn.size();
	spec.dSize = dConn.size();
}

bool startsAtZeroAndHasContinuity(const set<size_t>& sset) {
	size_t cntr = 0;
	set<size_t>::const_iterator setIter;
	for (setIter = sset.begin(); setIter != sset.end(); setIter++) {
		if (*setIter != cntr) {
			return false;
		}
		cntr++;
	}
	
	return true;
}

bool setOverlap(const set<size_t>& set1, const set<size_t>& set2) {
	
	set<size_t>::iterator findIter;
	findIter = find_first_of(set1.begin(), set1.end(), 
							 set2.begin(), set2.end());
	return (findIter != set1.end());
}

template<typename _value, typename _key>
const vector<_value> getValueVector(const map<_key, _value>& kvMap) {
	vector<_value> valueVec;
	transform(kvMap.begin(), kvMap.end(), valueVec.begin(), getValue<_value, _key>);
	return valueVec;
}

template<typename _value, typename _key>
const set<_key> getKeySet(const map<_key, _value>& kvMap) {
	vector<_key> keyVec;
	transform(kvMap.begin(), kvMap.end(), keyVec.begin(), getKey<_value, _key>);
	return set<_key>(keyVec.begin(), keyVec.end());
}

template<typename _value, typename _key>
const _value getValue(const pair<_key, _value>& kvPair) {
	return kvPair->second;
}

template<typename _value, typename _key>
const _key getKey(const pair<_key, _value>& kvPair) {
	return kvPair.first;
}
